# dataset_generator.py
"""
负责处理文档并生成用于微调的 RAG 数据集。
"""
import logging
from typing import List

from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.llama_dataset.generator import RagDatasetGenerator
from llama_index.core.llms import LLM
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import Document

async def generate_rag_dataset(
    documents: List[Document],
    llm: LLM,
    chunk_size: int,
    chunk_overlap: int,
    questions_per_chunk: int
) :
    """
    通过 LlamaIndex 管道将文档转换为 RAG 微调数据集。

    Args:
        documents (List[Document]): 原始文档列表。
        llm (LLM): 用于生成问题的语言模型。
        chunk_size (int): 句子分割的块大小。
        chunk_overlap (int): 块之间的重叠大小。
        questions_per_chunk (int): 每个块生成的问题数量。

    Returns:
        RagDataset: 生成的 RAG 数据集。
    """
    logging.info("开始处理文档并生成节点...")
    pipeline = IngestionPipeline(
        transformations=[
            SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap),
        ]
    )
    nodes = await pipeline.arun(documents=documents)
    logging.info(f"文档处理完成，生成了 {len(nodes)} 个节点。")

    logging.info("开始生成 RAG 数据集...")
    dataset_generator = RagDatasetGenerator(
        nodes,
        llm=llm,
        num_questions_per_chunk=questions_per_chunk,
        show_progress=True,
    )
    rag_dataset = await dataset_generator.agenerate_dataset_from_nodes()
    logging.info("RAG 数据集生成完毕。")
    return rag_dataset